package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"text/template"

	"gitee.com/bjf-fhe/merge/platform"
	"github.com/sirupsen/logrus"
	"github.com/tealeg/xlsx"

	_ "gitee.com/bjf-fhe/merge/platform/llama"
	_ "gitee.com/bjf-fhe/merge/platform/qianfan"
)

func main() {

	var sourceDir string
	var destinationFile string
	var instructionId, inputId, outputId int
	var outputTemplate string
	var jsonl bool
	var modelType string
	//parse flag to read the source dir
	flag.StringVar(&sourceDir, "source", "./datas", "source directory")
	//parse flag to read the destination dir
	flag.StringVar(&destinationFile, "dest", "./ft.json", "destination file")
	flag.StringVar(&outputTemplate, "output-template", "./output.templ", "output template file")
	flag.IntVar(&instructionId, "instruction-id", 1, "instruction column id")
	flag.IntVar(&inputId, "input-id", 2, "input column id")
	flag.IntVar(&outputId, "output-id", 3, "output column id")
	flag.BoolVar(&jsonl, "jsonl", false, "output jsonl format")
	flag.StringVar(&modelType, "model-type", "llama", "model type")
	flag.Parse()

	if jsonl {
		// if dest end with .json, change it to .jsonl
		if strings.HasSuffix(destinationFile, ".json") {
			destinationFile = strings.TrimSuffix(destinationFile, ".json") + ".jsonl"
		}
	}

	//检查outputTemplate是否存在，如果存在，则读取模板文件
	outputTempl, _ := template.ParseFiles(outputTemplate)

	//read all the csv in source dir
	var entries []platform.Entry
	var err = filepath.WalkDir(sourceDir, func(path string, d fs.DirEntry, err error) error {
		logrus.Infoln("read from file", path)
		var filename = d.Name()
		if strings.HasSuffix(filename, ".xlsx") && !strings.HasPrefix(filename, "~$") {
			// 打开文件
			xlFile, err := xlsx.OpenFile(path)
			if err != nil {
				fmt.Println(err.Error())
				return err
			}
			// 遍历sheet页读取
			for _, sheet := range xlFile.Sheets {
				//遍历行读取
				for _, row := range sheet.Rows {
					var entry = platform.NewEntry(modelType)
					// 遍历每行的列读取
					entry.SetInstruction(row.Cells[instructionId-1].String())
					if inputId > 0 && inputId <= len(row.Cells) {
						entry.SetInput(row.Cells[inputId-1].String())
						if outputTempl != nil {
							var w strings.Builder
							var cells = make(map[string]string)
							for k, cell := range row.Cells {
								cells[fmt.Sprintf("cell%d", k+1)] = cell.String()
							}
							err = outputTempl.Execute(&w, cells)
							if err == nil {
								output := w.String()
								if runtime.GOOS == "windows" {
									output = strings.ReplaceAll(output, "\r\n", "\n")
								}
								entry.SetOutput(output)
							}
						} else {
							entry.SetOutput(row.Cells[outputId-1].String())
						}
						entries = append(entries, entry)
					} else {
						logrus.WithField("row", row).WithField("cells", row.Cells).Warn("invalid row")
					}
				}
			}
			logrus.Infoln("import success")
		}
		return nil
	})
	if err == nil {
		var file *os.File
		//open the destination file and write entries as json
		file, err = os.Create(destinationFile)
		if err != nil {
			panic(err)
		}
		defer file.Close()
		if jsonl {
			for _, entry := range entries {
				err = json.NewEncoder(file).Encode(entry)
				if err != nil {
					break
				}
			}
		} else {
			logrus.WithField("entries", len(entries)).Info("entries")
			err = json.NewEncoder(file).Encode(entries)
		}
	}
	if err != nil {
		panic(err)
	}

}
